package com.chenwy.example.learning.trie;

class TrieNode{

    private TrieNode [] nodes;

    private boolean isLeaf;

    private char character;

    private TrieNode prev;

    private int size;

    public TrieNode(){
        this.nodes = new TrieNode[26];
        this.size = 0;
    }

    public boolean containsKey(char character){
        return this.nodes[character - 'a'] != null;
    }

    public void put(char character,TrieNode node){
        if(!containsKey(character)){
            node.character = character;
            this.nodes[character-'a'] = node;
            size ++;
        }
    }

    public TrieNode get(char character){
        return nodes[character - 'a'];
    }

    public void setIsLeaf(boolean isLeaf){
        this.isLeaf = isLeaf;
    }


    public boolean isLeaf(){
        return this.isLeaf;
    }

    public void setPrev(TrieNode prev){
        this.prev = prev;
    }

    public TrieNode getPrev(){
        return this.prev;
    }


    public void remove(char character){
        nodes[character - 'a'] = null;
        size --;//同时size减少
    }

    public void output(){
        for(int i = 0 ; i < nodes.length ;i ++){
            if(i != 0){
                System.out.print(" ");
            }
            if(nodes[i] != null){
                System.out.print((char)(i + 'a'));
            }else{
                System.out.print("null");
            }
        }
        System.out.println();
        System.out.println("size=" + size);
    }

    public char getCharacter(){
        return this.character;
    }

    public int size(){
        return this.size;
    }

    public boolean isEmpty(){
        return this.size == 0;
    }
}
public class Trie {

    private TrieNode root;

    private int size;

    public Trie(){
        this.root = new TrieNode();
        this.size = 0;
    }

    public void insert(String word){
        TrieNode node = root;
        for(int i = 0 ;i < word.length();i++){
            char character = word.charAt(i);
            if(!node.containsKey(character)){
                TrieNode n = new TrieNode();
                n.setPrev(node);
                node.put(character,n);
            }
            node = node.get(character);
        }
        if(!node.isLeaf()){
            node.setIsLeaf(true);
            size ++;
        }
    }

    public TrieNode searchPrefix(String word){
        TrieNode node = root;
        for(int i = 0 ;i < word.length(); i++){
            char character = word.charAt(i);
            if(!node.containsKey(character)){
                return null;
            }
            node = node.get(character);
        }
        return node;
    }

    public boolean search(String word){
        TrieNode node = searchPrefix(word);
        return node != null && node.isLeaf();
    }

    public boolean startsWith(String word){
        TrieNode node = searchPrefix(word);
        return node != null;
    }


    public void remove(String word){
        TrieNode node = searchPrefix(word);
        if(node != null && node.isLeaf()){
            if(!node.isEmpty() ){
                node.setIsLeaf(false);
            }else{
                while(node != null && node.getPrev() != null){
                    char character = node.getCharacter();
                    //删除后继节点
                    node.getPrev().remove(character);
                    //前驱节点是否是叶节点  或者 前驱存在其他后继节点
                    if(node.getPrev().isLeaf() || !node.getPrev().isEmpty()){
                        break;
                    }
                    node = node.getPrev();
                }
            }
            size -- ;
        }
    }

    public int size(){
        return this.size;
    }

    public static void main(String[] args) {
        Trie trie = new Trie();
        trie.insert("a");
        trie.insert("ab");
        trie.insert("c");
        trie.insert("d");
        trie.insert("e");
        trie.insert("f");
        trie.remove("a");
        System.out.println(trie.search("a"));
        System.out.println(trie.search("ab"));
        System.out.println(trie.size());
    }
}
